#!/bin/python
# -*- coding:utf-8 -*- 

import os,sys
import math
import numpy as np
import torch

AAindex=['GLY','ALA','VAL','LEU','ILE','PHE','TRP','TYR','ASP','ASN','GLU','LYS','GLN','MET','SER','THR','CYS','PRO','HIS','ARG']
aaindex='GAVLIFWYDNEKQMSTCPHR'
RNAindex=['A','U','C','G']
DNAindex = ['DA','DT','DC','DG']

def distance(c1,c2):
    n = 0
    for i in range(0,3):
        n += (c1[i] - c2[i]) ** 2
    return math.sqrt(n)

class Atom(list):
    def __init__(self,name,num,x,y,z):
        self.name = name
        self.num = num
        self.extend([x,y,z])

class Residue(list):
    def __init__(self,name,num,atoms):
        self.label = 0
        self.interface_label = 0
        self.name = name
        self.num = num
        self.extend(atoms)

class Chain(list):
    def __init__(self,name,residues):
        self.name = name
        self.extend(residues)

class Model(list):
    def __init__(self,name,chains):
        self.name = name
        self.extend(chains)

class ParseLine:
    def __init__(self,line):
        self.atom_name = line[12:16].strip()
        self.res_name = line[17:20].strip()
        self.chain_name = line[20:22].strip()
        self.atom_num = int(line[6:11].strip())
        self.res_num = int(line[22:26].strip())
        self.insert = line[26]
        self.x = float(line[30:38].strip())
        self.y = float(line[38:46].strip())
        self.z = float(line[46:54].strip())

def read_model(file_name):
    atoms = []
    residues = []
    chains = []
#	helix = []
    old_line = ''

    i = 0
    for line in open(file_name):
        if line[:6] == 'ENDMDL':
            break
        if (line[0:4] == 'ATOM'):
            line = ParseLine(line)
            if i > 0:
                if  line.res_num != old_line.res_num or line.res_name != old_line.res_name:
                    residues.append(Residue(old_line.res_name,old_line.res_num,atoms))
                    atoms = []
                if line.chain_name != old_line.chain_name:
                    chains.append(Chain(old_line.chain_name, residues))
                    residues = []

            if line.atom_name[0] == 'H':
                continue

            atoms.append(Atom(line.atom_name, line.atom_num, line.x, line.y, line.z))
            old_line = line
            i += 1

    residues.append(Residue(old_line.res_name, old_line.res_num, atoms))
    chains.append(Chain(old_line.chain_name, residues))
    return Model(file_name, chains)


def read_pdb(Model):
    for Chain in Model:
        coordN  = []
        coordCA = []
        coordC  = []
        coord  = []
        for Residue in Chain:
            #print(type(Residue.num))
            if Residue.num not in [63,64,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80]:
                continue
            for Atom in Residue:
                if Atom.name in ['N', 'CA', 'C']:
                    #vectorN = np.array([Atom[0],Atom[1],Atom[2]])                    
                    coord.append([Atom[0],Atom[1],Atom[2]])
                    #print(coord[-1])
                    #coordN.append([Atom[0],Atom[1],Atom[2]])
                    #print("coordN:", coordN)
                    
                #if Atom.name == 'CA':
                    #vectorCA=np.array([Atom[0],Atom[1],Atom[2]])                     
                #    coordCA.append([Atom[0],Atom[1],Atom[2]])

                #if Atom.name == 'C':
                    #vectorC=np.array([Atom[0],Atom[1],Atom[2]])                     
                #    coordC.append([Atom[0],Atom[1],Atom[2]])
                    
    return coord#coordN,coordCA,coordC 

'''
def main():
    A = read_model("./test.pdb")
    coord = read_pdb(A)
    #coordN,coordCA,coordC= read_pdb(A)
    #print(sys.argv[1])
    #print("coord:", coord)
    #print("coordN:", coordN)
    #print("coordCA:", coordCA)
    #print("coordC:", coordC)
    #extractChains(A)
    plan=[]
    #torch.tensor(coord[0])
    quats=[1,2,3,6]
    quat = torch.tensor(quats)
    for i in range(0, len(coord), 3):
        #plan.append(coord[i:i+3])
        torch.tensor(coord[i:i+3])
        torch.tensor(coord[i:i+3]).split
        print(torch.tensor(coord[i:i+3]).split(1))
        print( )
        print(torch.tensor(coord[i:i+3]).split(1)[0][:,1])
        print( )
    #    for ii in (coord[i:i+3]):
    #        torch.tensor(ii[0])
    #        print(torch.tensor(ii))
    #print(torch.tensor(coord[0]))
    #for coordi in coord:
    #    print(coordi)

if __name__ == '__main__':
    main()
'''
